//Michele Gregori 2009 miclegr@gmail.com

#include "dqDeformer.h"

MTypeId dqDeformer::id(0x8703);
MObject dqDeformer::matrix;
MObject dqDeformer::preBindMatrix;
MObject dqDeformer::geomMatrix;
MObject dqDeformer::lockWeights;
MObject dqDeformer::paintWeight;
MObject dqDeformer::paintTrans;
MObject dqDeformer::paintArrDirty;
MObject dqDeformer::bindPose;


void *dqDeformer::creator(){
	return new dqDeformer();
}

MStatus dqDeformer::initialize(){
	MFnMatrixAttribute mAttr;
	matrix=mAttr.create("matrix","mat");
	mAttr.setArray(1);
	mAttr.setReadable(0);
	mAttr.setKeyable(0); 
	mAttr.setConnectable(1);

	MFnMatrixAttribute preBAttr;
	preBindMatrix=preBAttr.create("bindPreMatrix","pbm");
	preBAttr.setArray(1);
	preBAttr.setReadable(0);
	preBAttr.setKeyable(0); 
	preBAttr.setConnectable(1);
	
	MFnMatrixAttribute geoMAttr;
	geomMatrix=geoMAttr.create("geomMatrix","gm");
	geoMAttr.setReadable(0);
	
	MFnNumericAttribute lockAttr;
	lockWeights=lockAttr.create("lockWeights","liw",MFnNumericData::kBoolean);
	lockAttr.setArray(1);
	
	MDoubleArray ar;
	MFnDoubleArrayData defVal;
	MFnTypedAttribute pwAttr;
	paintWeight=pwAttr.create("paintWeights", "ptw", MFnData::kDoubleArray);
	pwAttr.setReadable(0);
	pwAttr.setKeyable(0); 
	pwAttr.setConnectable(1);	
	pwAttr.setDefault(defVal.create(ar));
	
	MFnMessageAttribute mgAttr;
	paintTrans=mgAttr.create("paintTrans","ptt");
	mgAttr.setReadable(0);
	paintArrDirty=mgAttr.create("paintArrDirty","pad");
	mgAttr.setReadable(0);	
	bindPose=mgAttr.create("bindPose","bp");
	mgAttr.setReadable(0);
	
	addAttribute(matrix);
	addAttribute(preBindMatrix);
	addAttribute(geomMatrix);
	addAttribute(lockWeights);
	addAttribute(paintWeight);
	addAttribute(paintTrans);
	addAttribute(paintArrDirty);
	addAttribute(bindPose);

	return MS::kSuccess;
}

MStatus dqDeformer::shouldSave(const MPlug& inPlug, bool& outResult){
	if(inPlug==weights){
		outResult=true;
		return MStatus::kSuccess;
	}
	return MPxDeformerNode::shouldSave(inPlug, outResult);
}

void dqDeformer::postConstructor(){
	MPxDeformerNode::postConstructor();
	MThis=thisMObject();
	firstTime=true;
	inCreation=1;
}

MStatus dqDeformer::connectionMade(const MPlug &plug, const MPlug &otherPlug,bool asSrc){
	if(plug==inputGeom){
		MObject oldMesh=plug.asMObject();
		MFnMeshData dataFn;
		geomData=dataFn.create();
		meshFn=new MFnMesh(oldMesh);
		meshFn->copy(oldMesh,geomData);
		
		MPointArray pArray;
		meshFn->getPoints(pArray);
		initialMesh.resize(pArray.length());
		for(int x=0;x<pArray.length();x++){
			initialMesh.at(x)=pArray[x];
		}
		inCreation=0;
	}
	if(plug==matrix){
		
		MPlug geomMatrixPl(MThis,geomMatrix);
		MFnMatrixData mat(geomMatrixPl.asMObject());
		geomMatrixData=mat.matrix();
		
		int index=plug.logicalIndex();
		MPlug bindPl(MThis,preBindMatrix);
		MFnMatrixData dqMat(plug.asMObject());
		MFnMatrixData bindMat(bindPl.elementByLogicalIndex(index).asMObject());
		if(dqData.size()<=index){
			dqData.resize(index+1);
			bindData.resize(index+1);
		}
		bindData.at(index)=new MMatrix(bindMat.matrix());
		dqData.at(index)=new MDualQuaternion(geomMatrixData*(*bindData.at(index))*dqMat.matrix());
		
		MPlug weightListPl(MThis,weightList);
		MPlug weightPl(MThis,weights);
		int numVtx=weightPl.numElements();
		weightListData.resize(numVtx);
		for(int x=0;x<numVtx;x++){
			if(weightListData.at(x).size()<=index)weightListData.at(x).resize(index+1);
			double val=weightListPl.elementByPhysicalIndex(x).child(0).elementByLogicalIndex(index).asDouble();
			weightListData.at(x).at(index).weight=new double(val);
			weightListData.at(x).at(index).index=new unsigned int(x);
		}
	}
	if(plug==paintTrans){
		evalWeight=0;
		MPlugArray matCon;
		MPlug matrixListPlug(MThis,matrix);
		for(int x=0;x<matrixListPlug.numElements();x++){
			matrixListPlug[x].connectedTo(matCon,1,0);
			if(matCon[0].node()==otherPlug.node()){
				paintIndex=matrixListPlug[x].logicalIndex();
				break;
			}
		}
		MDoubleArray weightArray;
		weightArray.setLength(weightListData.size());
		for(iter iw=weightListData.begin();iw<weightListData.end();iw++){
			weightArray.set(*iw->at(paintIndex).weight,*iw->at(paintIndex).index);
		}
		
		MFnDoubleArrayData arrayData;
		MObject arrayObj=arrayData.create(weightArray);
		MPlug paintWeightPl(MThis,paintWeight);
		paintWeightPl.setValue(arrayObj);
	}
	if(plug==lockWeights){
		int index=plug.logicalIndex();
		bool val=plug.asBool();
		if(lockedJoint.size()<=index){
			lockedJoint.resize(index+1);
		}
		lockedJoint.at(index)=new bool(val);
		
	}
	
	return MPxDeformerNode::connectionMade( plug, otherPlug, asSrc );
}

MStatus dqDeformer::connectionBroken(const MPlug &plug, const MPlug &otherPlug,bool asSrc){
	if(plug==inputGeom){
		delete meshFn;
		for(int x=0;x<dqMesh.size();x++){
			delete dqMesh[x];
		}
	}
	if(plug==matrix){
		int index=plug.logicalIndex();
		delete dqData.at(index);
		delete bindData.at(index);
		for(int x=0;x<weightListData.size();x++){
			delete weightListData.at(x).at(index).weight;
			delete weightListData.at(x).at(index).index;
			
		}
	}
	if(plug==lockWeights){
		int index=plug.logicalIndex();
		delete lockedJoint.at(index);
		lockedJoint.at(index)=NULL;
	}
	return MPxDeformerNode::connectionBroken( plug, otherPlug, asSrc );
}


MStatus dqDeformer::compute(const MPlug &plug,MDataBlock &block){
	if(plug==outputGeom){
		if(firstTime){
			for(int x=0;x<dqMesh.size();x++){
				delete dqMesh[x];
			}
			dqMesh.resize(0);
			block.inputArrayValue(matrix);
			for(iter iw=weightListData.begin();iw<weightListData.end();iw++){
				MPoint pt=initialMesh[*iw->at(0).index];
				MDualQuaternion *dq=new MDualQuaternion;
				for(int y=0;y<iw->size();y++){
					*dq+=((*dqData[y])*(*iw->at(y).weight));					
				}
				dq->mulPoint(pt);
				dqMesh.push_back(dq);
				pt*=geomMatrixData.inverse();
				meshFn->setPoint(*iw->at(0).index,pt);
			}
			firstTime=false;
		}
		for (int x=dqDirty.size()-1;x>=0;x--){
			int matIndex=dqDirty[x];
			MArrayDataHandle arrayHnd=block.inputArrayValue(matrix);
			arrayHnd.jumpToArrayElement(matIndex);
			MDualQuaternion newDq(geomMatrixData*(*bindData.at(matIndex))*arrayHnd.inputValue().asMatrix());
			MPoint pt;
			for(iter iw=weightListData.begin();iw<weightListData.end();iw++){
				if(abs(*iw->at(matIndex).weight)<0.001)continue;
				pt=initialMesh[*iw->at(matIndex).index];
				*dqMesh[*iw->at(matIndex).index]-=(*dqData.at(matIndex))*(*iw->at(matIndex).weight);
				*iw->at(matIndex).weight=abs(*iw->at(matIndex).weight);
				if(newDq.dot(*dqMesh[*iw->at(matIndex).index])<0)(*iw->at(matIndex).weight)*=-1;
				*dqMesh[*iw->at(matIndex).index]+=newDq*(*iw->at(matIndex).weight);
				dqMesh[*iw->at(matIndex).index]->mulPoint(pt);
				pt*=geomMatrixData.inverse();
				meshFn->setPoint(*iw->at(matIndex).index,pt);
			}
			*dqData.at(matIndex)=newDq;
			dqDirty.pop_back();
		}
		meshFn->updateSurface();
		MDataHandle outHnd=block.outputValue(plug);
		outHnd.setMObject(geomData);			
		block.setClean(plug);
		return MS::kSuccess;
	}
	return MS::kUnknownParameter;
} 	

MStatus dqDeformer::setDependentsDirty(const MPlug &plugBeingDirtied,MPlugArray &affectedPlugs){
	if(inCreation)return MS::kSuccess;
	if(plugBeingDirtied==matrix){
		int index=plugBeingDirtied.logicalIndex();
		dqDirty.push_back(index);
		MPlug out(MThis,outputGeom);
		affectedPlugs.append(out.elementByPhysicalIndex(0));
		return MS::kSuccess;
	}
	if(plugBeingDirtied==paintWeight){
		firstTime=true;
		MFnDoubleArrayData arrayData(plugBeingDirtied.asMObject());
		MDoubleArray weightArray=arrayData.array();
		MPlug weightListPl(MThis,weightList);
		if(evalWeight){
			for(int x=0;x<weightArray.length();x++){
				if(weightArray[x]!=*weightListData[x][paintIndex].weight){
					int freeInf=weightListData[x].size();
					double totWeight=weightArray[x];
					double totLockWeight=(*lockedJoint[paintIndex] ? weightArray[x] : 0);
					int numInf=0;
					for(int y=0;y<weightListData[x].size();y++){
						if(y==paintIndex)continue;
						if(*weightListData[x][y].weight>0){
							numInf++;
							totWeight+=*weightListData[x][y].weight;
							if(*lockedJoint[y]) --freeInf,totLockWeight+=*weightListData[x][y].weight;
						}
					}
					double totFreeWeight=totWeight-totLockWeight-weightArray[x];
					if(!freeInf)continue;
					if(numInf>5){
					}
					if(totWeight!=1){
						int zeroW=0;
						double scaleFactor=(1-weightArray[x])/totFreeWeight;
						if((1-weightArray[x])==0)scaleFactor=0;
						for(int y=0;y<weightListData[x].size();y++){
							if(*lockedJoint[y])continue;
							if(y==paintIndex){
								*weightListData[x][y].weight=weightArray[x];
							}else{
								if(*weightListData[x][y].weight==0)zeroW++;
								*weightListData[x][y].weight*=scaleFactor;
							}
							weightListPl.elementByLogicalIndex(x).child(0).elementByLogicalIndex(y).setValue(*weightListData[x][y].weight);
						}
						if(zeroW==weightListData[x].size()-1){
							for(int y=0;y<weightListData[x].size();y++){
							if(y==paintIndex)continue;
							*weightListData[x][y].weight=(1-weightArray[x])/(weightListData[x].size()-1);
							weightListPl.elementByLogicalIndex(x).child(0).elementByLogicalIndex(y).setValue(*weightListData[x][y].weight);
							}
						}
					}
				}
			}
			MPlug out(MThis,outputGeom);
			affectedPlugs.append(out.elementByPhysicalIndex(0));
		}
		evalWeight=1;
		return MS::kSuccess;
	}
	if(plugBeingDirtied==lockWeights){
		int index=plugBeingDirtied.logicalIndex();
		bool val=plugBeingDirtied.asBool();
		if(lockedJoint.at(index)!=NULL){
			*lockedJoint.at(index)=(!val);
		}
		return MS::kSuccess;
	}
	return MS::kSuccess;
}
